{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c51e5438-dca6-4555-b953-e5247c1cdf1c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "===================================BUG REPORT===================================\n",
      "Welcome to bitsandbytes. For bug reports, please submit your error trace to: https://github.com/TimDettmers/bitsandbytes/issues\n",
      "================================================================================\n",
      "[2024-03-14 01:13:12,456] [INFO] [real_accelerator.py:161:get_accelerator] Setting ds_accelerator to cuda (auto detect)\n"
     ]
    }
   ],
   "source": [
    "\"\"\"\n",
    "init\n",
    "\"\"\"\n",
    "import clip\n",
    "import json\n",
    "import torch\n",
    "import faiss\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "from tqdm import tqdm\n",
    "Image.MAX_IMAGE_PIXELS = None\n",
    "\n",
    "### init clip\n",
    "clip_model, preprocess = clip.load(\"ViT-B/16\", device=\"cuda\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f15ebcf2",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "select the dataset and shot number\n",
    "dataset: caltech101, eurosat, dtd, food101, oxford_flowers, oxford_pets, stanford_cars, sun397, ucf101, ...\n",
    "shot number : 1, 2, 4, 8, 16\n",
    "\"\"\"\n",
    "shot_number = 4\n",
    "dataset_name = 'caltech101'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "8360c14d-a884-4862-97f7-ca6b0a55d5c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def index_save(index_split, path):\n",
    "    faiss.write_index(index_split, path)\n",
    "    return\n",
    "\n",
    "def get_number(str_data):\n",
    "    str_data = str_data.replace(\"'\", '\"')\n",
    "    data = json.loads(str_data)\n",
    "    label_value = data['label']\n",
    "    return label_value\n",
    "\n",
    "def find_positions_and_paths(gt_labels, image_paths, n):\n",
    "    positions = {}\n",
    "    for idx, label in enumerate(gt_labels):\n",
    "        if label not in positions:\n",
    "            positions[label] = [idx]\n",
    "        elif len(positions[label]) < n:\n",
    "            positions[label].append(idx)\n",
    "    return positions\n",
    "\n",
    "def write_to_file(str_list, num_list, file_path):\n",
    "    with open(file_path, 'w', encoding='utf-8') as file:\n",
    "        for s, n in zip(str_list, num_list):\n",
    "            file.write(f\"{s} {n}\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "63cb2fa5-82f2-4f30-8820-84834f0f0ccc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "47\n",
      "['banded', 'blotchy', 'braided', 'bubbly', 'bumpy', 'chequered', 'cobwebbed', 'cracked', 'crosshatched', 'crystalline', 'dotted', 'fibrous', 'flecked', 'freckled', 'frilly', 'gauzy', 'grid', 'grooved', 'honeycombed', 'interlaced', 'knitted', 'lacelike', 'lined', 'marbled', 'matted', 'meshed', 'paisley', 'perforated', 'pitted', 'pleated', 'polka-dotted', 'porous', 'potholed', 'scaly', 'smeared', 'spiralled', 'sprinkled', 'stained', 'stratified', 'striped', 'studded', 'swirly', 'veined', 'waffled', 'woven', 'wrinkled', 'zigzagged']\n"
     ]
    }
   ],
   "source": [
    "\"\"\"\n",
    "get the classnames\n",
    "\"\"\"\n",
    "classnames_file_path = f\"/mnt/petrelfs/liuziyu/LLM_Memory/SimplyRetrieve/CLIP-Cls/benchmarks_test/{dataset_name}_database/classnames.txt\"\n",
    "with open(classnames_file_path, 'r') as file:\n",
    "    classnames = file.readlines()\n",
    "print(len(classnames))\n",
    "classnames = [classname.strip() for classname in classnames]\n",
    "print(classnames)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec331e2d-3e6e-47dd-b1cb-53e751b764d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "### save_file_path stores the selected k-shot images and their order in a .txt file, preparing for retrieving later\n",
    "save_file_path = f\"./database/{dataset_name}_database/{dataset_name}_{shot_number}_shot_database.txt\"\n",
    "### file_path refer to the file which includes the whole image path and labels of trainset. you can get this file in CLIP_Cls fold\n",
    "file_path = f\"/mnt/petrelfs/liuziyu/LLM_Memory/SimplyRetrieve/CLIP-Cls/benchmarks_test/{dataset_name}_database/trainset.txt\"\n",
    "\n",
    "with open(file_path, 'r') as file:\n",
    "    lines = file.readlines()\n",
    "image_paths = [line.split()[0] for line in lines]\n",
    "\n",
    "labels = [line.split(' ',1)[1] for line in lines]\n",
    "gt_labels = [get_number(label) for label in labels]\n",
    "# labels = [line.split()[1] for line in lines]\n",
    "# gt_labels = [int(label) for label in labels]\n",
    "\n",
    "print(len(image_paths))\n",
    "print(len(gt_labels))\n",
    "data_postition = find_positions_and_paths(gt_labels,image_paths,shot_number)\n",
    "print(data_postition)\n",
    "\n",
    "### normal\n",
    "select_gt_labels = []\n",
    "select_image_path = []\n",
    "for i in range(len(classnames)):\n",
    "    image_index = data_postition[i]\n",
    "    for j in range(len(image_index)):\n",
    "        select_gt_labels.append(i)\n",
    "        select_image_path.append(image_paths[image_index[j]])\n",
    "print(select_gt_labels)\n",
    "print(select_image_path)\n",
    "write_to_file(select_image_path, select_gt_labels, save_file_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "4f5fac82-f18f-4740-9c12-507ff042475d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "752\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Process:: 100%|██████████| 752/752 [00:09<00:00, 76.32it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total number of indexes: 752\n",
      "Saving index:\n",
      "Done\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "\"\"\"\n",
    "build memory\n",
    "\"\"\"\n",
    "file_path = f\"./database/{dataset_name}_database/{dataset_name}_{shot_number}_shot_database.txt\"\n",
    "with open(file_path, 'r') as file:\n",
    "    lines = file.readlines()\n",
    "image_paths = [line.split()[0] for line in lines]\n",
    "print(len(image_paths))\n",
    "\n",
    "### index save path\n",
    "index_img_save_path =  f\"./database/{dataset_name}_database/{dataset_name}_{shot_number}_shot_img_index.index\"\n",
    "\n",
    "### build index\n",
    "index_img = faiss.IndexHNSWFlat(512, 64, faiss.METRIC_INNER_PRODUCT)\n",
    "\n",
    "### embedding images\n",
    "embed_img = []\n",
    "with torch.no_grad():\n",
    "    for image_path in tqdm(image_paths,desc=\"Process:\"):\n",
    "        image = preprocess(Image.open(image_path)).unsqueeze(0).to(\"cuda\")\n",
    "        image_features = clip_model.encode_image(image)\n",
    "        image_features /= image_features.norm(dim=-1, keepdim=True)\n",
    "        embed_img.append(image_features.cpu())\n",
    "embed_img = [np.array(embed) for embed in embed_img]\n",
    "embed_img = np.array(embed_img).squeeze()\n",
    "index_img.add(embed_img)\n",
    "print(\"Total number of indexes:\", index_img.ntotal)\n",
    "\n",
    "### save index\n",
    "print(\"Saving index:\")\n",
    "index_save(index_img, index_img_save_path)\n",
    "print(\"Done\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "7b71ee42-d58c-4e1a-9b42-ce766d9824d8",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Process:: 100%|██████████| 3783/3783 [00:42<00:00, 89.89it/s] \n"
     ]
    }
   ],
   "source": [
    "### get CLIP+KNN prediction results\n",
    "index_img_save_path = f\"./database/{dataset_name}_database/{dataset_name}_{shot_number}_shot_img_index.index\"\n",
    "trainset_file_path = f\"./database/{dataset_name}_database/{dataset_name}_{shot_number}_shot_database.txt\"\n",
    "predictions_save_path = f\"/mnt/petrelfs/liuziyu/LLM_Memory/SimplyRetrieve/CLIP-Cls/output/ZeroshotCLIP_topk/vit_b16/{dataset_name}/predictions_{shot_number}_shot_knn.pth\"\n",
    "index = faiss.read_index(index_img_save_path)\n",
    "pth_file_path = f'/mnt/petrelfs/liuziyu/LLM_Memory/SimplyRetrieve/CLIP-Cls/output/ZeroshotCLIP_topk/vit_b16/{dataset_name}/predictions.pth'\n",
    "predictions = torch.load(pth_file_path)\n",
    "\n",
    "with open(trainset_file_path, 'r') as file:\n",
    "    lines = file.readlines()\n",
    "    for prediction in tqdm(predictions,desc=\"Process:\"):\n",
    "        ### 解析pth文件，获取图片位置和原来的预测结果\n",
    "        for item in prediction.values():\n",
    "            pre_class = item['pred_class']\n",
    "            # print(item['label'])\n",
    "        for item in prediction.keys():\n",
    "            test_img_path = item\n",
    "    \n",
    "        with torch.no_grad():\n",
    "            image = preprocess(Image.open(test_img_path)).unsqueeze(0).to(\"cuda\")\n",
    "            # torch.Size([1, 512])\n",
    "            image_features = clip_model.encode_image(image)\n",
    "            image_features /= image_features.norm(dim=-1, keepdim=True)\n",
    "            image_features = np.array(image_features.cpu())\n",
    "            distance, index_result = index.search(image_features, 100)\n",
    "            \n",
    "            labels = []\n",
    "            for index_number in index_result[0]:\n",
    "                parts = lines[index_number].strip().split(' ', 1)\n",
    "                part1, part2 = parts\n",
    "                labels.append(int(part2))\n",
    "            labels = torch.tensor(labels)\n",
    "            # print(labels)\n",
    "        \n",
    "            ### 修改pth文件\n",
    "            for item in prediction.values():\n",
    "                item['pred_class'] = labels\n",
    "        \n",
    "torch.save(predictions, predictions_save_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e578c3c3-a5ce-4f7f-b733-8b4128706c02",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
